package com.zk.algorithm.tree;

import com.zk.algorithm.annotation.LeetCodeExplore;
import com.zk.algorithm.annotation.Medium;
import com.zk.algorithm.bean.TreeNode;

import java.util.Stack;

@LeetCodeExplore
@Medium
public class KthSmallestElementInABST {

    // ===================================
    // Kth smallest
    // 方法一，通过统计左右树节点的数量，确定当前位于哪个 key
    // ===================================

    public int kthSmallest(TreeNode root, int k) {
        int count = countNodes(root.left);
        if (k <= count) {
            return kthSmallest(root.left, k);
        } else if (k > count + 1) {
            return kthSmallest(root.right, k - 1 - count); // 1 is counted as current node
        }

        //      命中
        //       ↓
        return root.val;
    }

    public int countNodes(TreeNode n) {
        if (n == null) return 0;

        return 1 + countNodes(n.left) + countNodes(n.right);
    }

    // ===================================
    // 中序遍历
    // ===================================

    int count = 0;
    int result = Integer.MIN_VALUE;

    public int kthSmallest2(TreeNode root, int k) {
        traverse(root, k);
        return result;
    }

    public void traverse(TreeNode root, int k) {
        if (root == null) return;
        traverse(root.left, k);

        count++;
        if (count == k) result = root.val;

        traverse(root.right, k);
    }

    // ===================================
    // 中序迭代版本
    // ===================================

    public int kthSmallest3(TreeNode root, int k) {
        Stack<TreeNode> stack = new Stack<TreeNode>();
        TreeNode p = root;
        int count = 0;

        while (!stack.isEmpty() || p != null) {
            if (p != null) {
                stack.push(p);  // Just like recursion
                p = p.left;
            } else {
                TreeNode node = stack.pop();
                if (++count == k) return node.val;
                p = node.right;
            }
        }

        return Integer.MIN_VALUE;
    }

}
